import sys
import os
import asyncio
import pygame
from PyQt6.QtWidgets import (
    QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, QTextEdit,
    QPushButton, QLabel, QComboBox, QSlider, QFileDialog, QMessageBox, 
    QGroupBox, QLineEdit, QProgressBar
)
from PyQt6.QtCore import Qt, QThread, pyqtSignal
from PyQt6.QtGui import QIcon, QFont
from tts_engine import TTSEngine

# 初始化pygame用于音频播放
pygame.mixer.init()

class AsyncWorker(QThread):
    """异步任务工作线程"""
    finished = pyqtSignal(bool, str)
    progress = pyqtSignal(int)
    
    def __init__(self, task, *args, **kwargs):
        super().__init__()
        self.task = task
        self.args = args
        self.kwargs = kwargs
        
    def run(self):
        try:
            # 使用事件循环运行异步任务
            loop = asyncio.new_event_loop()
            asyncio.set_event_loop(loop)
            result = loop.run_until_complete(self.task(*self.args, **self.kwargs))
            loop.close()
            self.finished.emit(*result)
        except Exception as e:
            self.finished.emit(False, str(e))

class TTSApp(QMainWindow):
    def __init__(self):
        super().__init__()
        self.tts_engine = TTSEngine()
        self.init_ui()
        self.setWindowTitle("Edge-TTS 文本转语音工具")
        self.setGeometry(100, 100, 800, 600)
        
        # 应用图标
        if hasattr(sys, '_MEIPASS'):
            icon_path = os.path.join(sys._MEIPASS, 'assets', 'icon.png')
        else:
            icon_path = 'assets/icon.png'
        self.setWindowIcon(QIcon(icon_path))
        
        # 当前播放文件
        self.current_audio_file = ""
        
    def init_ui(self):
        """初始化用户界面"""
        main_widget = QWidget()
        main_layout = QVBoxLayout()
        
        # 标题
        title_label = QLabel("Edge-TTS 文本转语音工具")
        title_label.setFont(QFont("Arial", 16, QFont.Weight.Bold))
        title_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        main_layout.addWidget(title_label)
        
        # 语音设置区域
        settings_group = QGroupBox("语音设置")
        settings_layout = QVBoxLayout()
        
        # 语音选择
        voice_layout = QHBoxLayout()
        voice_layout.addWidget(QLabel("选择语音:"))
        self.voice_combo = QComboBox()
        
        # 获取中文语音名称列表
        voice_names = self.tts_engine.get_voice_names()
        self.voice_combo.addItems(voice_names)
        
        # 查找并选中"云健(男)"语音
        for i, name in enumerate(voice_names):
            if "云健" in name:
                self.voice_combo.setCurrentIndex(i)
                break
        
        self.voice_combo.currentIndexChanged.connect(self.voice_changed)
        voice_layout.addWidget(self.voice_combo, 3)
        
        # 语速设置
        rate_layout = QHBoxLayout()
        rate_layout.addWidget(QLabel("语速:"))
        self.rate_slider = QSlider(Qt.Orientation.Horizontal)
        self.rate_slider.setRange(-50, 50)
        self.rate_slider.setValue(0)
        self.rate_slider.setTickPosition(QSlider.TickPosition.TicksBelow)
        self.rate_slider.setTickInterval(10)
        rate_layout.addWidget(self.rate_slider, 3)
        self.rate_label = QLabel("+0%")
        rate_layout.addWidget(self.rate_label)
        
        # 音量设置
        volume_layout = QHBoxLayout()
        volume_layout.addWidget(QLabel("音量:"))
        self.volume_slider = QSlider(Qt.Orientation.Horizontal)
        self.volume_slider.setRange(-50, 50)
        self.volume_slider.setValue(0)
        self.volume_slider.setTickPosition(QSlider.TickPosition.TicksBelow)
        self.volume_slider.setTickInterval(10)
        volume_layout.addWidget(self.volume_slider, 3)
        self.volume_label = QLabel("+0%")
        volume_layout.addWidget(self.volume_label)
        
        # 代理设置
        proxy_layout = QHBoxLayout()
        proxy_layout.addWidget(QLabel("代理设置:"))
        self.proxy_edit = QLineEdit()
        self.proxy_edit.setPlaceholderText("http://proxy_ip:port (可选)")
        proxy_layout.addWidget(self.proxy_edit)
        
        # 添加到设置区域
        settings_layout.addLayout(voice_layout)
        settings_layout.addLayout(rate_layout)
        settings_layout.addLayout(volume_layout)
        settings_layout.addLayout(proxy_layout)  # 添加代理设置
        settings_group.setLayout(settings_layout)
        main_layout.addWidget(settings_group)
        
        # 文本输入区域
        text_group = QGroupBox("输入文本")
        text_layout = QVBoxLayout()
        self.text_edit = QTextEdit()
        self.text_edit.setPlaceholderText("请输入要转换为语音的文本...")
        
        # 连接文本变化信号
        self.text_edit.textChanged.connect(self.update_char_count)
        
        text_layout.addWidget(self.text_edit)
        
        # 添加字符计数显示
        char_count_layout = QHBoxLayout()
        self.char_count_label = QLabel("字符数: 0")
        char_count_layout.addWidget(self.char_count_label)
        
        # 添加最大字符数提示
        self.max_char_label = QLabel("最大支持 5000 字符")
        char_count_layout.addWidget(self.max_char_label, alignment=Qt.AlignmentFlag.AlignRight)
        
        text_layout.addLayout(char_count_layout)
        
        text_group.setLayout(text_layout)
        main_layout.addWidget(text_group)
        
        # 按钮区域
        button_layout = QHBoxLayout()
        
        self.generate_btn = QPushButton("生成并播放语音")
        self.generate_btn.clicked.connect(self.generate_speech)
        button_layout.addWidget(self.generate_btn)
        
        self.stop_btn = QPushButton("停止播放")
        self.stop_btn.clicked.connect(self.stop_playback)
        self.stop_btn.setEnabled(False)
        button_layout.addWidget(self.stop_btn)
        
        self.batch_btn = QPushButton("批量处理文件")
        self.batch_btn.clicked.connect(self.batch_process)
        button_layout.addWidget(self.batch_btn)
        
        main_layout.addLayout(button_layout)
        
        # 输出设置
        output_layout = QHBoxLayout()
        output_layout.addWidget(QLabel("输出文件名:"))
        self.filename_edit = QLineEdit("output/output.mp3")
        output_layout.addWidget(self.filename_edit)
        
        self.browse_btn = QPushButton("浏览...")
        self.browse_btn.clicked.connect(self.browse_output)
        output_layout.addWidget(self.browse_btn)
        
        main_layout.addLayout(output_layout)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setRange(0, 100)
        self.progress_bar.setValue(0)
        self.progress_bar.setVisible(False)
        main_layout.addWidget(self.progress_bar)
        
        # 状态栏
        self.status_label = QLabel("就绪")
        self.status_label.setAlignment(Qt.AlignmentFlag.AlignCenter)
        main_layout.addWidget(self.status_label)
        
        # 连接滑块值改变信号
        self.rate_slider.valueChanged.connect(self.update_rate_label)
        self.volume_slider.valueChanged.connect(self.update_volume_label)
        
        # 初始更新标签
        self.update_rate_label(0)
        self.update_volume_label(0)
        
        main_widget.setLayout(main_layout)
        self.setCentralWidget(main_widget)
    
    def update_char_count(self):
        """更新字符计数"""
        text = self.text_edit.toPlainText()
        char_count = len(text)
        self.char_count_label.setText(f"字符数: {char_count}")
        
        # 如果超过5000字符，显示警告
        if char_count > 5000:
            self.char_count_label.setStyleSheet("color: red;")
        else:
            self.char_count_label.setStyleSheet("")
    
    def voice_changed(self, index):
        """语音选择改变"""
        self.tts_engine.set_voice(index)
    
    def update_rate_label(self, value):
        """更新语速标签"""
        sign = '+' if value >= 0 else ''
        self.rate_label.setText(f"{sign}{value}%")
    
    def update_volume_label(self, value):
        """更新音量标签"""
        sign = '+' if value >= 0 else ''
        self.volume_label.setText(f"{sign}{value}%")
    
    def generate_speech(self):
        """生成语音文件并播放"""
        text = self.text_edit.toPlainText().strip()
        if not text:
            QMessageBox.warning(self, "输入错误", "请输入要转换的文本！")
            return
        
        # 检查字符数
        char_count = len(text)
        if char_count > 5000:
            QMessageBox.warning(self, "文本过长", "文本超过5000字符，请缩短文本！")
            return
        
        output_file = self.filename_edit.text().strip()
        if not output_file:
            QMessageBox.warning(self, "输出错误", "请输入输出文件名！")
            return
        
        # 确保文件扩展名正确
        if not output_file.lower().endswith(('.mp3', '.wav')):
            output_file += '.mp3'
        
        self.status_label.setText("正在生成语音文件...")
        self.progress_bar.setVisible(True)
        self.progress_bar.setValue(30)
        
        # 确保生成的字符串始终包含符号
        rate_sign = '+' if self.rate_slider.value() >= 0 else ''
        volume_sign = '+' if self.volume_slider.value() >= 0 else ''
        
        rate = f"{rate_sign}{self.rate_slider.value()}%"
        volume = f"{volume_sign}{self.volume_slider.value()}%"
        
        # 获取代理设置
        proxy = self.proxy_edit.text().strip() or None
        
        # 创建并启动工作线程
        self.worker = AsyncWorker(
            self.tts_engine.text_to_speech,
            text, output_file, rate, volume, proxy
        )
        self.worker.finished.connect(self.on_generation_finished)
        self.worker.start()
    
    def on_generation_finished(self, success, result):
        """生成完成回调"""
        self.progress_bar.setValue(100)
        
        if success:
            self.status_label.setText(f"文件已保存: {self.filename_edit.text()}")
            
            # 播放生成的音频文件
            self.play_audio(self.filename_edit.text())
            
            QMessageBox.information(self, "完成", "语音文件生成成功！")
        else:
            QMessageBox.critical(self, "生成错误", f"无法生成语音文件: {result}")
            self.status_label.setText("生成失败")
        
        self.progress_bar.setVisible(False)
    
    def play_audio(self, file_path):
        """播放音频文件"""
        try:
            # 停止当前播放
            pygame.mixer.music.stop()
            
            # 加载并播放新文件
            pygame.mixer.music.load(file_path)
            pygame.mixer.music.play()
            
            self.current_audio_file = file_path
            self.stop_btn.setEnabled(True)
            self.status_label.setText("正在播放音频...")
        except Exception as e:
            QMessageBox.critical(self, "播放错误", f"无法播放音频: {str(e)}")
            self.status_label.setText("播放失败")
    
    def stop_playback(self):
        """停止播放"""
        pygame.mixer.music.stop()
        self.stop_btn.setEnabled(False)
        self.status_label.setText("播放已停止")
    
    def browse_output(self):
        """浏览输出文件位置"""
        file_path, _ = QFileDialog.getSaveFileName(
            self,
            "保存语音文件",
            self.filename_edit.text(),
            "音频文件 (*.mp3 *.wav)"
        )
        if file_path:
            self.filename_edit.setText(file_path)
    
    def batch_process(self):
        """批量处理文件"""
        files, _ = QFileDialog.getOpenFileNames(
            self,
            "选择文本文件",
            "",
            "文本文件 (*.txt)"
        )
        if not files:
            return
        
        self.status_label.setText(f"准备批量处理 {len(files)} 个文件...")
        self.progress_bar.setVisible(True)
        self.progress_bar.setRange(0, len(files))
        self.progress_bar.setValue(0)
        
        # 获取输出目录
        output_dir = QFileDialog.getExistingDirectory(
            self,
            "选择输出目录"
        )
        if not output_dir:
            self.status_label.setText("批量处理已取消")
            self.progress_bar.setVisible(False)
            return
        
        # 确保生成的字符串始终包含符号
        rate_sign = '+' if self.rate_slider.value() >= 0 else ''
        volume_sign = '+' if self.volume_slider.value() >= 0 else ''
        
        rate = f"{rate_sign}{self.rate_slider.value()}%"
        volume = f"{volume_sign}{self.volume_slider.value()}%"
        
        # 获取代理设置
        proxy = self.proxy_edit.text().strip() or None
        
        # 处理每个文件
        for i, file_path in enumerate(files):
            self.progress_bar.setValue(i)
            self.status_label.setText(f"正在处理: {os.path.basename(file_path)} ({i+1}/{len(files)})")
            QApplication.processEvents()  # 更新UI
            
            # 读取文本内容
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    text = f.read().strip()
                
                if not text:
                    continue
                
                # 检查字符数
                char_count = len(text)
                if char_count > 5000:
                    self.status_label.setText(f"跳过: {os.path.basename(file_path)} (超过5000字符)")
                    continue
                
                # 生成输出文件名
                base_name = os.path.splitext(os.path.basename(file_path))[0]
                output_file = os.path.join(output_dir, f"{base_name}.mp3")
                
                # 生成语音文件
                loop = asyncio.new_event_loop()
                asyncio.set_event_loop(loop)
                success, error = loop.run_until_complete(
                    self.tts_engine.text_to_speech(text, output_file, rate, volume, proxy))
                loop.close()
                
                if not success:
                    self.status_label.setText(f"处理失败: {os.path.basename(file_path)}")
                    QMessageBox.warning(self, "处理失败", f"无法处理文件 {os.path.basename(file_path)}: {error}")
            
            except Exception as e:
                self.status_label.setText(f"处理失败: {os.path.basename(file_path)}")
                QMessageBox.warning(self, "处理失败", f"处理文件 {os.path.basename(file_path)} 时出错: {str(e)}")
        
        self.progress_bar.setValue(len(files))
        self.status_label.setText(f"批量处理完成: {len(files)} 个文件")
        QMessageBox.information(self, "完成", f"成功处理 {len(files)} 个文件！")
        self.progress_bar.setVisible(False)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    
    # 设置应用样式
    app.setStyle("Fusion")
    
    window = TTSApp()
    window.show()
    sys.exit(app.exec())